from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading

# analysis of phase space

def progress_bar(current, total, barLength = 100):
    percent = float(current) * 100 / total
    arrow = '-' * int(percent/100 * barLength - 1) + '>'
    spaces = ' ' * (barLength - len(arrow))

    print('Progress: [%s%s] %d %%' % (arrow, spaces, percent), end='\r')
    sys.stdout.flush()

@partial(jax.jit, static_argnums=(3,))
def model_forward(variables, state, x, model):
    """ forward pass of the model
    """
    return model.apply(variables, state, x)

@jit
def get_action(y):
    return jnp.argmax(y)
get_action_vmap = jax.vmap(get_action)

# load landscape and states from file
def load_task(pth = "./logs/task.json", display = True):
    # open json file
    with open(pth, "r") as f:
        data = json.load(f)
        landscape = data["data"]
        state = data["state"]
        goal = data["goal"]
        if display:
            print("state: ", state)
            print("goal: ", goal)
            print("landscape: ", landscape)
    return landscape, state, goal

# 全局变量，用于存储和图像显示线程交互的数据
class imgview:
    global_image = None
    imgview_exit = False
    trajectory = []
    focus_i = 0
    traj_i = 0

imgview_data = imgview()

# 定义一个函数，用于在独立线程中显示图像
def show_image():
    grid_size_display = 20
    while not imgview_data.imgview_exit:
        # 检查全局变量是否有图像
        if imgview_data.global_image is not None:
            img = np.copy(imgview_data.global_image)
            state_x = imgview_data.trajectory[imgview_data.traj_i][0]
            state_y = imgview_data.trajectory[imgview_data.traj_i][1]
            cv2.circle(img, (state_y * grid_size_display + int(grid_size_display/2), state_x * grid_size_display + int(grid_size_display/2)), 7, (0, 0, 255), -1, cv2.LINE_AA)
            # 显示图像
            cv2.imshow("Image", img)
            key = cv2.waitKey(1)
            if key == ord('a'):
                imgview_data.focus_i -= 1
                print("imgview_data.focus_i: ", imgview_data.focus_i)
            elif key == ord('d'):
                imgview_data.focus_i += 1
                print("imgview_data.focus_i: ", imgview_data.focus_i)
        else:
            # 图像还未产生，等待100毫秒
            time.sleep(0.1)

def ivf():

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)
    parser.add_argument("--start_i", type=int, default=rpl_config.start_i)
    parser.add_argument("--end_i", type=int, default=rpl_config.end_i)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration
    rpl_config.start_i = args.start_i
    rpl_config.end_i = args.end_i

    """ load model
    """
    params = load_weights(rpl_config.model_pth)
    # 定义一个函数，用于生成随机权重
    def init_weights_r(key, shape):
        return jax.random.normal(key, shape)
    # 生成一个随机的 PRNGKey
    key = jax.random.PRNGKey(np.random.randint(0, 1000))
    # key = jax.random.PRNGKey(4512)
    # 使用 tree_map 遍历 params 对象，并使用 init_weights 函数生成随机权重
    random_params = jax.tree_map(lambda x: init_weights_r(key, x.shape), params)
    # params = random_params

    # get elements of params
    tree_leaves = jax.tree_util.tree_leaves(params)
    for i in range(len(tree_leaves)):
        print("shape of leaf ", i, ": ", tree_leaves[i].shape)

    bias1 = np.array(tree_leaves[0])
    mat1 = np.array(tree_leaves[1])
    print("mat1.shape: ", mat1.shape)
    # print("bias1: ", bias1)

    mat_obs = np.array(tree_leaves[1])[rpl_config.nn_size:rpl_config.nn_size+9,:]
    mat_intr = np.array(tree_leaves[1])[0:rpl_config.nn_size,:]
    print("mat_obs.shape: ", mat_obs.shape)
    
    """ create agent
    """
    if rpl_config.nn_type == "vanilla":
        model = RNN(hidden_dims = rpl_config.nn_size)
    elif rpl_config.nn_type == "gru":
        model = GRU(hidden_dims = rpl_config.nn_size)

    # check if param fits the agent
    if rpl_config.nn_type == "vanilla":
        assert params["params"]["Dense_0"]["kernel"].shape[0] == rpl_config.nn_size + 10

    n_samples = 10000
    k1 = npr.randint(0, 1000000)
    rnn_state = model.initial_state_rnd(n_samples, k1)
    rnn_state_old = rnn_state.copy()
    diff = jnp.abs(rnn_state - rnn_state_old)
    rnn_state_old = rnn_state.copy()
    diff_norm = jnp.linalg.norm(diff, axis=1)
    diff_norm_old = diff_norm.copy()
    norm_std = diff_norm.copy()

    rnn_state_init = rnn_state.copy()

    obs_zero = jnp.array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0] for i in range(n_samples)])

    rnn_state_trajectory = []
    intr_field = []

    for t in range(rpl_config.life_duration):

        progress_bar(t, rpl_config.life_duration)

        intr_vector = np.dot(rnn_state, mat_intr) + bias1
        intr_vector = np.tanh(intr_vector)
        intr_field.append(intr_vector - rnn_state)

        rnn_state_trajectory.append(np.array(rnn_state).copy())
        
        """ model forward 
        """
        # rnn_state, y1 = model_forward(params, rnn_state, obs_zero, model)
        rnn_state = intr_vector
        
    print(rnn_state.shape)
    print(norm_std.shape)
    rnn_state_np = np.array(rnn_state)

    # 将 rnn_state_trajectory 展开成 rnn_state_np 的形状
    rnn_state_trajectory_np = np.array(rnn_state_trajectory)
    rnn_state_trajectory_np = rnn_state_trajectory_np.reshape(-1, rnn_state_trajectory_np.shape[-1])
    print("shape of rnn_state_trajectory_np: ", rnn_state_trajectory_np.shape)

    # 将 intr_field 展开成 rnn_state_np 的形状
    intr_field_np = np.array(intr_field)
    intr_field_np = intr_field_np.reshape(-1, intr_field_np.shape[-1])
    print("shape of intr_field_np: ", intr_field_np.shape)

    # 对 rnn_state_np 进行 PCA
    intrinsic_pca = PCA()
    # pca.fit(rnn_state_np)
    intrinsic_pca.fit(rnn_state_trajectory_np)

    # # 打印 variance ratio
    # print(pca.explained_variance_ratio_)

    rnn_state_trajectory_np_pca = intrinsic_pca.transform(rnn_state_trajectory_np)

    print("shape of rnn_state_trajectory_np_pca: ", rnn_state_trajectory_np_pca.shape)

    intr_pca = intrinsic_pca.transform(intr_field_np + intrinsic_pca.mean_)

    compute_dim0 = 0
    compute_dim1 = 1

    @jax.jit
    def compute_intr_field(HS_trajectory_1, mat_intr, bias1):
        intr_vector = jnp.dot(HS_trajectory_1, mat_intr) + bias1
        intr_vector = jnp.tanh(intr_vector)
        intr_field = intr_vector - HS_trajectory_1
        return intr_field
    compute_intr_field_vmap = vmap(compute_intr_field, in_axes=(0, None, None))

    def compute_diff_rnd_sample(compute_dim0, compute_dim1):

        compute_dims = [compute_dim0]

        # 将除了指定维度以外的所有维度置零
        mask = np.zeros(rnn_state_trajectory_np_pca.shape[1], dtype=bool)
        mask[compute_dims] = False
        rnn_state_trajectory_np_pca_masked = rnn_state_trajectory_np_pca.copy()
        start_time = time.time()
        rnn_state_trajectory_np_pca_masked[:, ~mask] = 0
        end_time = time.time()
        # print("time elapsed6: ", end_time - start_time)

        # # 将 pca 空间的点，逆变换回原始相空间
        start_time = time.time()
        HS_trajectory_1 = intrinsic_pca.inverse_transform(rnn_state_trajectory_np_pca_masked)
        end_time = time.time()
        # print("time elapsed5: ", end_time - start_time)

        # 计算主成分空间中的内禀向量场
        start_time = time.time()
        intr_field1 = compute_intr_field_vmap(HS_trajectory_1, mat_intr, bias1)
        end_time = time.time()
        # print("time elapsed4: ", end_time - start_time)

        compute_dims = [compute_dim0, compute_dim1]

        # 将除了指定维度以外的所有维度置零
        start_time = time.time()
        mask = np.zeros(rnn_state_trajectory_np_pca.shape[1], dtype=bool)
        mask[compute_dims] = True
        rnn_state_trajectory_np_pca_masked = rnn_state_trajectory_np_pca.copy()
        rnn_state_trajectory_np_pca_masked[:, ~mask] = 0
        end_time = time.time()
        # print("time elapsed3: ", end_time - start_time)

        # # 将 pca 空间的点，逆变换回原始相空间
        start_time = time.time()
        HS_trajectory_1 = intrinsic_pca.inverse_transform(rnn_state_trajectory_np_pca_masked)
        end_time = time.time()
        # print("time elapsed2: ", end_time - start_time)
        
        # 计算主成分空间中的内禀向量场
        start_time = time.time()
        intr_field2 = compute_intr_field_vmap(HS_trajectory_1, mat_intr, bias1)
        end_time = time.time()
        # print("time elapsed1: ", end_time - start_time)
        
        start_time = time.time()
        diff_ = np.array(intr_field1) - np.array(intr_field2)
        diff_norm = np.linalg.norm(diff_, axis=1)
        diff_norm_mean = np.mean(diff_norm)
        end_time = time.time()
        # print("time elapsed0: ", end_time - start_time)
        
        return diff_norm_mean
    
    # # monitor start time
    # start_time = time.time()
    # diff_norm_mean = compute_diff_rnd_sample(compute_dim0, compute_dim1)
    # # monitor end time
    # end_time = time.time()
    # print("diff_norm_mean: ", diff_norm_mean)
    # print("time elapsed: ", end_time - start_time)

    # create a 32x32 matrix
    dim_corr_mat = np.zeros((rpl_config.nn_size, rpl_config.nn_size))
    # fill the matrix with diff_norm_mean
    for i in range(rpl_config.nn_size):
        for j in range(rpl_config.nn_size):
            start_time = time.time()
            dim_corr_mat[i, j] = compute_diff_rnd_sample(i, j)
            end_time = time.time()
            print("time elapsed: ", end_time - start_time)
            progress_bar(i * rpl_config.nn_size + j, rpl_config.nn_size * rpl_config.nn_size)

    print("dim_corr_mat: ", dim_corr_mat)
    np.save("./logs/dim_corr_mat.npy", dim_corr_mat)


if __name__ == "__main__":
    
    ivf()